#include"system_ws.h"
#include<fstream>
#include<complex>
using std::string;
using std::ifstream;
using std::complex;

typedef complex<double> complexd;

void System_Tz_WS::setupOrbitals()
{
  ifstream forb(OrbFile.c_str());
  int points;
  forb>>points;
  vector<double> r(points);
  vector<double> w(points);
  for(int i=0;i<points;i++)
    {
      forb>>r[i]>>w[i];
    }
  int numOrbitals;
  forb>>numOrbitals;
  Orbitals.resize(2*numOrbitals);
  int num,type,l,jj;
  double ereal,eimag,kreal,kimag;
  for(int i=0;i<numOrbitals;i++)
    {
      forb>>num>>type>>l>>jj>>ereal>>eimag>>kreal>>kimag;
      ereal*=hbar_c*hbar_c*0.5/mc2;
      eimag*=hbar_c*hbar_c*0.5/mc2;
      Orbitals[2*i].set(type,l,jj,-1,complexd(ereal,eimag),complexd(kreal,kimag));
      Orbitals[2*i].r=r;
      Orbitals[2*i].w=w;
      Orbitals[2*i].wavefun.resize(points);
      Orbitals[2*i].up.resize(points);      
      Orbitals[2*i].V.resize(points);
      double wavefun_real,wavefun_imag,up_real,up_imag;
      for(int j=0;j<points;j++)
	{
	  forb>>wavefun_real>>wavefun_imag>>up_real>>up_imag;
	  Orbitals[2*i].wavefun[j]=complexd(wavefun_real,wavefun_imag);
	  Orbitals[2*i].up[j]=complexd(up_real,up_imag);
	  forb>>Orbitals[2*i].V[j];
	}
      //copy to neutron
      Orbitals[2*i+1]=Orbitals[2*i];
      Orbitals[2*i+1].tz=1;
    }
  setupGroups();

  OneBodyMat.resize(Groups.size());
  for(int i=0;i<Groups.size();i++)
    {
      int dim=Groups[i].size();
      OneBodyMat[i].resize(dim*(dim+1)/2);
    }
  
  ifstream f1B(OneBodyFile.c_str());
  int lines;  
  f1B>>lines;
  for(int i=0;i<lines;i++)
    {
      int a,b;
      double k2real;
      double k2imag;
      complexd k2;
      f1B>>a>>b>>k2real>>k2imag;
      k2=complexd(k2real,k2imag);
      a--;//start from zero
      b--;
      int ap=2*a;
      int bp=2*b;
      int group=GroupIndice[ap];//for proton
      int aap=OrbIndice[ap];
      int bbp=OrbIndice[bp];
      int dim=Groups[group].size();
      OneBodyMat[group][mapab_e(aap,bbp,dim)]=k2*hbar_c*hbar_c*0.5/mc2;

      int an=2*a+1;
      int bn=2*b+1;
      group=GroupIndice[an];//for neutron
      int aan=OrbIndice[an];
      int bbn=OrbIndice[bn];      
      dim=Groups[group].size();
      OneBodyMat[group][mapab_e(aan,bbn,dim)]=k2*hbar_c*hbar_c*0.5/mc2;
    }

  setupOverlap();
}

/*
  dim of MatEle=4,the 4 double are:
      0  interaction
      1  (p1 \dot p2)/(m hbar omega) + m omega r1 \dot r2 / hbar
      2  (r1-r2)^2 m omega/(2 hbar)
      3  (p1 \dot p2)/(m hbar omega)
*/
const int System_Tz_WS::MatEledim=System_TzType::MatEledim;
void System_Tz_WS::setupTwoBodyMat()
{
  setupIndiceIn2BStates();
  int totalChannels=TwoBodyChannels.size();
  TwoBodyMat.resize(totalChannels);
  #pragma omp parallel for
  for(int channel=0;channel<totalChannels;channel++)
    {
      int num2BInChannel=TwoBodyChannels[channel].size();
      TwoBodyMat[channel].clear();
      TwoBodyMat[channel].resize( (num2BInChannel*(num2BInChannel+1))/2 );
      for(int bra=0;bra<num2BInChannel;bra++)
	for(int ket=bra;ket<num2BInChannel;ket++)
	  {
	    int ab=TwoBodyChannels[channel][bra];
	    int cd=TwoBodyChannels[channel][ket];
	    int a=TwoBodyStates[ab].a;
	    int b=TwoBodyStates[ab].b;
	    int c=TwoBodyStates[cd].a;
	    int d=TwoBodyStates[cd].b;
	    int J=TwoBodyStates[ab].J;
	    TwoBodyMatEle MatEle(MatEledim);
	    get2BmatOnNeed(a,b,c,d,J,MatEle);
	    set2Bmat(channel,bra,ket,MatEle);
	  }
    }
}


 void System_Tz_WS::get2BmatOnNeed(int a,int b,int c,int d,int J,TwoBodyMatEle & MatEle) const
 {
  MatEle.resize(MatEledim);
  MatEle.setZero();
  int jja=Orbitals[a].jj;
  int jjb=Orbitals[b].jj;
  int jjc=Orbitals[c].jj;
  int jjd=Orbitals[d].jj;
  int Tzab=(Orbitals[a].tz+Orbitals[b].tz)/2;
  int Tzcd=(Orbitals[c].tz+Orbitals[d].tz)/2;

  if(a==b && J%2) return;
  if(c==d && J%2) return;
  if(Tzab!=Tzcd) return;//isospin projection check
  if((Orbitals[a].l+Orbitals[b].l+Orbitals[c].l+Orbitals[d].l)%2) return;//parity check
  int Jmin=max(abs(jja-jjb)/2,abs(jjc-jjd)/2);
  int Jmax=min((jja+jjb)/2,(jjc+jjd)/2);
  if(Jmin>Jmax) return;//angular momenta check

  for(int ia=0;ia<Orbitals[a].Coeff.size();ia++)
    for(int ib=0;ib<Orbitals[b].Coeff.size();ib++)
      for(int ic=0;ic<Orbitals[c].Coeff.size();ic++)
	for(int id=0;id<Orbitals[d].Coeff.size();id++)
	  {
	    System_TzType::TwoBodyMatEle Temp;
	    int iia=pSystem_Tz->Groups[Orbitals[a].GroupIndex][ia];
	    int iib=pSystem_Tz->Groups[Orbitals[b].GroupIndex][ib];
	    int iic=pSystem_Tz->Groups[Orbitals[c].GroupIndex][ic];
	    int iid=pSystem_Tz->Groups[Orbitals[d].GroupIndex][id];
	    pSystem_Tz->get2BmatAt(iia,iib,iic,iid,J,Temp);
	    if(Temp.empty()) continue;
	    if(iia==iib) Temp*=sqrt(2.);
	    if(iic==iid) Temp*=sqrt(2.);
	    MatEle+=TwoBodyMatElement<complexd>(Temp)*(Orbitals[a].Coeff[ia]*Orbitals[b].Coeff[ib]*Orbitals[c].Coeff[ic]*Orbitals[d].Coeff[id]);
	  }
  if(a==b) MatEle*=sqrt(2.)/2;
  if(c==d) MatEle*=sqrt(2.)/2;
}

void System_Tz_WS::setupOverlap()
{
  double lengthb=pSystem_Tz->lengthb;
  int totalGroups=pSystem_Tz->Groups.size();
  int num=Orbitals.size();
  for(int i=0;i<num;i++)
    {
      int l=Orbitals[i].l;
      int jj=Orbitals[i].jj;
      int tz=Orbitals[i].tz;
      int group=0;//find the group of l,jj,tz
      for(group=0;group<totalGroups;group++)
	{
	  const System_TzType::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][0] ];
	  if(HOOrbital.l==l && HOOrbital.jj==jj && HOOrbital.tz==tz)
	    break;
	}
      if(group>=totalGroups)
	{
	  std::cout<<"warning: below channel in WS orbitals not found in HO. orbitals, set to zero:\n";
	  std::cout<<l<<"\t"<<jj<<"\t"<<tz<<std::endl;

	  Orbitals[i].GroupIndex=0;
	  int dimho=pSystem_Tz->Groups[0].size();	  
	  Orbitals[i].Coeff.resize(dimho);
	  for(int j=0;j<dimho;j++)//every HO. orbitals in a group
	    {
	      Orbitals[i].Coeff[j]=0;
	    }	  
	}
      else
	{
	  Orbitals[i].GroupIndex=group;
	  int dimho=pSystem_Tz->Groups[group].size();

	  Orbitals[i].Coeff.resize(dimho);
	  for(int j=0;j<dimho;j++)//every HO. orbitals in a group
	    {
	      const System_Tz_HO::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][j] ];
	      int n=HOOrbital.n;	  
	      DataType temp=0;//overlap with the HO. orbital
	      for(int k=0;k<Orbitals[i].r.size();k++)
		{
		  double r=Orbitals[i].r[k];
		  double w=Orbitals[i].w[k];
		  temp+=Orbitals[i].wavefun[k]*ho_r(n,l,lengthb,r)*r*w;
		}
	      Orbitals[i].Coeff[j]=temp;
	    }
	}
    }
}

complexd System_Tz_WS::get1B(int A,int bra,int ket) const
{
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]))
    {
      int group=GroupIndice[bra];
      int a=OrbIndice[bra];
      int b=OrbIndice[ket];
      int dim=Groups[group].size();
      if(a>b) swap(a,b);
      return (1-1./A)*OneBodyMat[group][mapab_e(a,b,dim)];
    }
  else
    return 0;
}

complexd System_Tz_WS::get1B_p2(int bra,int ket) const
{
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]))
    {
      int group=GroupIndice[bra];
      int a=OrbIndice[bra];
      int b=OrbIndice[ket];
      int dim=Groups[group].size();
      if(a>b) swap(a,b);
      return OneBodyMat[group][mapab_e(a,b,dim)];
    }
  else
    {
      return 0;
    }
}

void System_Tz_WS::printOrbitals(int num) const
{
  if(num==-1)
    num=Orbitals.size();
  std::cout<<"num"<<"\t"<<"l"<<"\t"<<"jj"<<"\t"<<"tz"<<"\t"<<"e"<<"\t"<<"type"<<endl;      
  for(int i=0;i<num;i++)
    {
      std::cout<<i<<"\t"<<Orbitals[i].l<<"\t"<<Orbitals[i].jj<<"\t"<<Orbitals[i].tz<<"\t"<<Orbitals[i].e<<"\t"<<Orbitals[i].type<<"\t"<<abs(InOverOut(i))<<std::endl;
    }
}

complexd System_Tz_WS::InOverOut(int orb) const
{
  int last=Orbitals[orb].r.size()-1;
  double r=Orbitals[orb].r[last];
  complexd u=Orbitals[orb].wavefun[last];
  complexd up=Orbitals[orb].up[last];
  complexd k=Orbitals[orb].k;
  int l=Orbitals[orb].l;
  complexd z=k*r;
  complexd jn,yn,djn,dyn;
  csphbessel_jy(l,z,jn,yn,djn,dyn);
  complexd ci(0,1.);
  complexd out=(jn+ci*yn)*z;
  complexd in=(jn-ci*yn)*z;
  complexd outp=(jn+z*djn+ci*(yn+z*dyn))*k;
  complexd inp=(jn+z*djn-ci*(yn+z*dyn))*k;
  //   /  out    in  \   / a \     / u \
  //  |               | |     | = |     |
  //   \  outp   inp /   \ b /     \ up/

  complexd bOvera=(out*up-outp*u)/(inp*u-in*up);
  return bOvera;
} 
